import numpy as np
import torch
import lfads
import wiener
from parameters import get_lfads_params, get_wiener_params


class Nomad(object):
    """
    :param data: 所用的神经信号数据，一个列表，列表中单个数据大小为(data_length,num_neurons)
    :param labels: 所用的标签数据，一个列表，列表中单个数据大小为(data_length,num_dimension)
    :param split_size: split_size取值0-1，用于划分数据集
    """

    def __init__(self, data, labels, split_size):
        self.data = data
        self.labels = labels
        self.split_size = split_size

    def update_lfads_data(self, data_loader, lfads_params):
        """
        :param data_loader: dataloader的形式的数据集,里面单个数据维度是三维的（1,data_length,num_neurons）,标签维度是二维的（data_length,num_dimension）
        :param lfads_params: 训练lfads模型所需要的参数
        """
        # 训练lfads模型
        aligner = lfads.train_lfads(data_loader, lfads_params)

        # 获取参数
        input_size = lfads_params['input_size']
        hidden_size = lfads_params['hidden_size']
        output_size = lfads_params['output_size']
        factor_size = lfads_params['factor_size']
        controller_size = lfads_params['controller_size']

        # 初始化训练好的模型LFADS
        model = lfads.LFADS(input_size=input_size, hidden_size=hidden_size, output_size=output_size,
                            factor_size=factor_size, controller_size=controller_size)
        model.load_state_dict(aligner)

        # 获取新的数据
        new_data = []
        new_label = []
        # 将数据集通过模型输出新的数据
        for i in range(len(self.data)):
            new_label.append(self.labels[i])
            lfads_data = torch.tensor(self.data[i]).float()

            recon_x, kl_loss = model(lfads_data)

            lfads_data = recon_x.detach().numpy()
            lfads_data = np.squeeze(lfads_data)
            new_data.append(lfads_data)

        # 更新数据集
        self.data = new_data
        self.labels = new_label

    def wiener_fit(self, wiener_params):
        """
        :param wiener_params: 训练wiener滤波器所需要的参数
        """
        # 获取数据集长度
        num_trial = len(self.data)

        # 划分数据集
        split_num_trail = int(num_trial * self.split_size)
        train_data = self.data[:split_num_trail]
        train_labels = self.labels[:split_num_trail]
        test_data = self.data[split_num_trail:]
        test_labels = self.labels[split_num_trail:]

        # 训练wiener滤波器
        wiener.train_wiener(train_data, train_labels, test_data, test_labels, wiener_params)


if __name__ == '__main__':
    # data为一个列表，列表中单个数据大小为(data_length,num_neurons)
    data = 'data'
    # labels为一个列表，列表中单个数据大小为(data_length,num_dimension)
    labels = 'labels'
    # 输入为dataloader的形式的数据集,里面单个数据维度是三维的（1,data_length,num_neurons）,标签维度是二维的（data_length,num_dimension）
    data_loader = 'dataloader'

    # 获取lfads所需要的参数
    lfads_params = get_lfads_params()
    # 获取wiener所需要的参数
    wiener_params = get_wiener_params()

    # 构建Nomad
    model = Nomad(data, labels, split_size=0.7)
    # 训练lfads模型并使用训练好的模型获取新的数据
    model.update_lfads_data(data_loader, lfads_params)
    # 使用新的数据进行wiener滤波训练与测试
    model.wiener_fit(wiener_params)
